import torch
import torch.nn as nn

# 定义卷积块的类，包含两个卷积层（每个卷积层后面跟一个BatchNorm层和ReLU激活函数层）
class ConvBlock(nn.Module):
    def __init__(self, ch_in, ch_out):
        super(ConvBlock, self).__init__()
        # 使用 nn.Sequential 封装卷积、BatchNorm和ReLU激活函数层，构建一个卷积块
        self.conv = nn.Sequential(
            nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=1, padding=1, bias=True),  # 卷积层：输入通道数 ch_in，输出通道数 ch_out，3x3 卷积核，填充为 1
            nn.BatchNorm2d(ch_out),  # Batch Normalization，对输出进行标准化处理
            nn.ReLU(inplace=True),  # ReLU激活函数，增加网络的非线性特性
            nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1, bias=True),
            nn.BatchNorm2d(ch_out),
            nn.ReLU(inplace=True))

    def forward(self, x):
        x = self.conv(x)  # 将输入 x 通过卷积块处理，得到卷积特征
        return x

# 定义上采样块的类，包含一个上采样层（双线性插值上采样）和一个卷积块
class UpConvBlock(nn.Module):
    def __init__(self, ch_in, ch_out):
        super().__init__()
        # 使用 nn.Sequential 封装上采样层和卷积层，构建一个上采样块
        self.up = nn.Sequential(
            nn.Upsample(scale_factor=2),  # 上采样层，双线性插值上采样，将特征图尺寸扩大两倍
            nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=1, padding=1, bias=True),
            nn.BatchNorm2d(ch_out),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.up(x)   # 将输入 x 通过上采样块处理，得到上采样特征
        return x


class UNet(nn.Module):
    def __init__(self, ch_in=3, ch_out=1):
        super().__init__()
        feature_channels = [8, 16, 32, 64, 128]   # 定义特征通道数列表

        # 定义四个池化层，池化核尺寸均为2，步长均为2,用于下采样
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        # 根据通道设定，定义下采样过程中使用的五个卷积块,作为编码器部分
        self.Conv1 = ConvBlock(ch_in, feature_channels[0])
        self.Conv2 = ConvBlock(feature_channels[0], feature_channels[1])
        self.Conv3 = ConvBlock(feature_channels[1], feature_channels[2])
        self.Conv4 = ConvBlock(feature_channels[2], feature_channels[3])
        self.Conv5 = ConvBlock(feature_channels[3], feature_channels[4])

        # 根据通道设定，定义上次样过程中使用的卷积层和上采样层,作为解码器部分
        self.Up5 = UpConvBlock(feature_channels[4], feature_channels[3])
        self.UpConv5 = ConvBlock(feature_channels[4], feature_channels[3])
        self.Up4 = UpConvBlock(feature_channels[3], feature_channels[2])
        self.UpConv4 = ConvBlock(feature_channels[3], feature_channels[2])
        self.Up3 = UpConvBlock(feature_channels[2], feature_channels[1])
        self.UpConv3 = ConvBlock(feature_channels[2], feature_channels[1])
        self.Up2 = UpConvBlock(feature_channels[1], feature_channels[0])
        self.UpConv2 = ConvBlock(feature_channels[1], feature_channels[0])

        # 定义最后一层卷积，输出通道数等于目标类别数
        self.Conv = nn.Conv2d(feature_channels[0], ch_out, kernel_size=1, stride=1, padding=0)
        self.sigmoid = nn.Sigmoid()  # Sigmoid激活函数，将输出限制在0到1之间

    def forward(self, x):
        # 第一层卷积，将3通道的输入变换为特征图
        f1 = self.Conv1(x)

        # 计算下采样过程，通过池化和卷积操作逐渐减小特征图的尺寸和通道数,获得中间特征图
        f2 = self.pool1(f1)
        f2 = self.Conv2(f2)
        f3 = self.pool2(f2)
        f3 = self.Conv3(f3)
        f4 = self.pool3(f3)
        f4 = self.Conv4(f4)
        f5 = self.pool4(f4)
        f5 = self.Conv5(f5)

        # 特征融合阶段：逐步将编码器的特征图与解码器阶段的特征图进行融合
        # 第一次特征融合
        up_f5 = self.Up5(f5)  # 上采样 f5
        up_f5 = torch.cat((f4, up_f5), dim=1)  # 拼接特征图 f4 和上采样后的 f5
        up_f5 = self.UpConv5(up_f5)  # 卷积融合特征

        # 第二次特征融合
        up_f4 = self.Up4(up_f5)  # 上采样 f4
        up_f4 = torch.cat((f3, up_f4), dim=1)  # 拼接特征图 f3 和上采样后的 f4
        up_f4 = self.UpConv4(up_f4)  # 卷积融合特征

        # 第三次特征融合
        up_f3 = self.Up3(up_f4)  # 上采样 f3
        up_f3 = torch.cat((f2, up_f3), dim=1)  # 拼接特征图 f2 和上采样后的 f3
        up_f3 = self.UpConv3(up_f3)  # 卷积融合特征

        # 第四次特征融合
        up_f2 = self.Up2(up_f3)  # 上采样 f2
        up_f2 = torch.cat((f1, up_f2), dim=1)   # 拼接特征图 f1 和上采样后的 f2
        up_f2 = self.UpConv2(up_f2)  # 卷积融合特征

        # 计算最后一层卷积输出，获得预测mask
        mask = self.Conv(up_f2)
        mask = self.sigmoid(mask)  # 将输出限制在 0 到 1 之间，获得预测 mask
        return mask


if __name__ == "__main__":
    unet = UNet()
    x = torch.randn(size=(1, 3, 256, 256))  # 构造一个形状为 (1, 3, 256, 256) 的张量作为输入
    print(unet(x).size())  # 打印模型输出的尺寸
